import torch
from torch import nn


class NumberNet(nn.Module):
    def __init__(self, device=None, classes=10):
        super().__init__()
        self.classes = classes
        self.device = device
        if device is None:  # 如果device没有定义
            # 判断device是否启用GPU训练
            self.device = torch.device("cpu")
            if torch.cuda.is_available():
                self.device = torch.device("cuda:0")

        # 输入层
        self.conv1 = nn.Conv2d(3, 32, 3, device=self.device)
        # 卷积层（此处卷积之后的大小为10）
        self.conv2 = nn.Conv2d(32, 64, 4, device=self.device)
        # 最大池化（将图片进行大小倍率缩减）
        self.pool = nn.MaxPool2d(2, 2)
        # 将最终的卷积图像进行像素拉伸
        self.flatten = nn.Flatten()
        # 设计线性回归进行分类
        self.fc1 = nn.Linear(1600, 1000, device=self.device)
        self.fc2 = nn.Linear(1000, self.classes, device=self.device)
        # 设计激活函数
        self.relu = nn.ReLU()
        self.softmax = nn.LogSoftmax(dim=-1)  # 输出
        # 削减神经元数量
        self.dropout = nn.Dropout()

    def forward(self, X):  # 28x28x3
        X = self.relu(self.conv1(X))  # 26x26x32
        X = self.pool(X)  # 13x13x32
        X = self.relu(self.conv2(X))  # 10x10x64
        X = self.pool(X)  # 5x5x64
        X = self.flatten(X)  # 1600x1
        X = self.dropout(X)
        X = self.relu(self.fc1(X))
        X = self.dropout(X)
        X = self.relu(self.fc2(X))
        return self.softmax(X)
